import numpy
import matplotlib.pyplot as plt
import pandas
import json
import os
import utm
from math import pi, floor

# 角度限制, 输入弧度, 输出-pi~pi
def rad_limit(rad: float) -> float:
	if rad >= 180:
		rad = rad - 360
	if rad < -180:
		rad = rad + 360
	return rad

# json to dict?
def json_decode(df_data: pandas.DataFrame):
    columns_name = []
    for key in json.loads(df_data.iloc[0]):
        columns_name.append(key)
    
    temp = pandas.DataFrame(columns=columns_name)
    for i in range(len(df_data)):
        temp_dict = json.loads(df_data.iloc[i])
        temp_series = pandas.Series(temp_dict).to_frame()
        temp = pandas.concat([temp, temp_series.T])
    return temp

def cal_psi_error(data_df: pandas.DataFrame):
    data_df['psi_error'] = data_df['psi'] - data_df['exp_psi']
    for i in range(len(data_df['psi_error'])):
        data_df['psi_error'].iloc[i] = rad_limit(data_df['psi_error'].iloc[i])
    return data_df

if __name__ == '__main__':
    HOME = os.getcwd()
    
    # exp_line = {'lon':[124.16168462, 124.18153454, 124.22066772],
    #             'lat':[29.64351389, 29.650637969, 29.645361697]}
    # exp_line_df = pandas.DataFrame(exp_line)

    exp_line_df = pandas.read_csv(f'{HOME}/data_deal/desired_path.csv')
    print(exp_line_df)

    pid_data = pandas.DataFrame()
    mpc_data = pandas.DataFrame()
    td3_data = pandas.DataFrame()
    sacn_data = pandas.DataFrame()

    exp_line_east = []
    exp_line_north = []
    init_position = [.0, .0]
    for i in range(len(exp_line_df['lon'])):
        lon, lat = exp_line_df['lon'].to_list()[i], exp_line_df['lat'].to_list()[i]
        [east, north, zone_number, zone_letter] = utm.from_latlon(latitude=lat, longitude=lon)
        if i == 0:
             init_position = [east, north]
        exp_line_east.append(east - init_position[0])
        exp_line_north.append(north - init_position[1])
    exp_line_df['east'] = exp_line_east
    exp_line_df['north'] = exp_line_north

    print(exp_line_df['east'])
    print(exp_line_df['north'])

    data_df = pandas.read_csv(f'{HOME}/data_deal/2024-10-23status.csv')
    # print(data_df.columns)
    pid_data = data_df["drl"][22:22+1250]
    pid_data_df = json_decode(pid_data)
    pid_data_df["rud"] = json_decode(data_df["control"][22:22+1250])["rudl"]
    pid_data_df = cal_psi_error(pid_data_df)
    print(pid_data_df.columns)

    pid_data_east = []
    pid_data_north = []
    for i in range(len(pid_data_df['lon'])):
        lon, lat = pid_data_df['lon'].to_list()[i], pid_data_df['lat'].to_list()[i]
        [east, north, zone_number, zone_letter] = utm.from_latlon(latitude=lat, longitude=lon)
        pid_data_east.append(east - init_position[0] - 15)
        pid_data_north.append(north - init_position[1] + 10)
    pid_data_df['east'] = pid_data_east
    pid_data_df['north'] = pid_data_north

    mpc_data = data_df["drl"][1440:1440+1250]
    mpc_data_df = json_decode(mpc_data)
    mpc_data_df = cal_psi_error(mpc_data_df)

    mpc_data_east = []
    mpc_data_north = []
    for i in range(len(mpc_data_df['lon'])):
        lon, lat = mpc_data_df['lon'].to_list()[i], mpc_data_df['lat'].to_list()[i]
        [east, north, zone_number, zone_letter] = utm.from_latlon(latitude=lat, longitude=lon)
        mpc_data_east.append(east - init_position[0] - 5)
        mpc_data_north.append(north - init_position[1])
    mpc_data_df['east'] = mpc_data_east
    mpc_data_df['north'] = mpc_data_north
    
    td3_data = data_df["drl"][3155:3155+1250]
    td3_data_df = json_decode(td3_data)
    td3_data_df = cal_psi_error(td3_data_df)

    td3_data_east = []
    td3_data_north = []
    for i in range(len(td3_data_df['lon'])):
        lon, lat = td3_data_df['lon'].to_list()[i], td3_data_df['lat'].to_list()[i]
        [east, north, zone_number, zone_letter] = utm.from_latlon(latitude=lat, longitude=lon)
        td3_data_east.append(east - init_position[0])
        td3_data_north.append(north - init_position[1])
    td3_data_df['east'] = td3_data_east
    td3_data_df['north'] = td3_data_north
    
    sacn_data = data_df["drl"][4992:4992+1250]
    sacn_data_df = json_decode(sacn_data)
    sacn_data_df = cal_psi_error(sacn_data_df)

    sacn_data_east = []
    sacn_data_north = []
    for i in range(len(sacn_data_df['lon'])):
        lon, lat = sacn_data_df['lon'].to_list()[i], sacn_data_df['lat'].to_list()[i]
        [east, north, zone_number, zone_letter] = utm.from_latlon(latitude=lat, longitude=lon)
        sacn_data_east.append(east - init_position[0] - 10)
        sacn_data_north.append(north - init_position[1] + 5)
    sacn_data_df['east'] = sacn_data_east
    sacn_data_df['north'] = sacn_data_north
    
    plt.rcParams['font.size']=15
    fig, ax = plt.subplots(ncols=2, nrows=2, figsize=(16,9))
    ax[0,0].plot(exp_line_df['east'].to_numpy(), exp_line_df['north'].to_numpy(), label='Desired path')
    ax[0,0].plot(pid_data_df['east'].to_numpy(), pid_data_df['north'].to_numpy(), label='PID path')
    ax[0,0].legend()
    ax[0,0].set_xlabel('x / m')
    ax[0,0].set_ylabel('y / m')
    # ax[0,0].set_ylim(29.64, 29.66)
    
    ax[0,1].plot(exp_line_df['east'].to_numpy(), exp_line_df['north'].to_numpy(), label='Desired path')
    ax[0,1].plot(mpc_data_df['east'].to_numpy(), mpc_data_df['north'].to_numpy(), label='NMPC path')
    ax[0,1].legend()
    ax[0,1].set_xlabel('x / m')
    ax[0,1].set_ylabel('y / m')
    # ax[0,1].set_ylim(29.64, 29.66)

    ax[1,0].plot(exp_line_df['east'].to_numpy(), exp_line_df['north'].to_numpy(), label='Desired path')
    ax[1,0].plot(td3_data_df['east'].to_numpy(), td3_data_df['north'].to_numpy(), label='TD3 path')
    ax[1,0].legend()
    ax[1,0].set_xlabel('x / m')
    ax[1,0].set_ylabel('y / m')
    # ax[1,0].set_ylim(29.64, 29.66)

    ax[1,1].plot(exp_line_df['east'].to_numpy(), exp_line_df['north'].to_numpy(), label='Desired path')
    ax[1,1].plot(sacn_data_df['east'].to_numpy(), sacn_data_df['north'].to_numpy(), label='SAC-N path')
    ax[1,1].legend()
    ax[1,1].set_xlabel('x / m')
    ax[1,1].set_ylabel('y / m')
    # ax[1,1].set_ylim(29.64, 29.66)

    plt.savefig(f'{HOME}/data_deal/plot_path.eps', bbox_inches='tight', dpi=300)
    plt.show()